/**
 * 文 件 名:  FileSystemClassLoader.java
 * 工 程 名:  ScriptManager
 * 创建日期:  2015年3月2日 下午4:08:28
 * 创建作者:  杨 强  <281455776@qq.com>
 */
package loader;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.util.Collection;
import java.util.HashSet;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 自定义文件加载器
 * 
 * @author 杨 强
 *
 */

public class FileSystemClassLoader extends ClassLoader {
    private static final Logger LOGGER = LoggerFactory.getLogger(FileSystemClassLoader.class);
    /**
     * 根目录
     */
    private String rootDir;
    /**
     * 打包的脚本jar文件
     */
    private String scriptJarFileName;
    /**
     * 需要由该类加载器直接加载的类
     */
    private HashSet<String> directLoadClassSet = new HashSet<String>();

    /**
     * @param rootDir
     *            脚本根路径
     * @param scriptJarFileName
     *            打包的脚本jar文件
     */
    public FileSystemClassLoader(String rootDir, String scriptJarFileName) {
        this.rootDir = rootDir;
        this.scriptJarFileName = scriptJarFileName;
    }

    /**
     * 清空直接加载的类
     */
    public void clearDirectLoadClass() {
        directLoadClassSet.clear();
    }

    /**
     * 添加由该加载器直接加载的类
     * 
     * @param className
     * @return
     */
    public FileSystemClassLoader addDirectLoadClass(String className) {
        directLoadClassSet.add(className);
        return this;
    }

    /**
     * 添加由该加载器直接加载的类
     * 
     * @param cls
     * @return
     */
    public FileSystemClassLoader addDirectLoadClass(Class<?> cls) {
        return addDirectLoadClass(cls.getName());
    }

    /**
     * 添加由该加载器直接加载的类的集合
     * 
     * @param clss
     * @return
     */
    public <T> FileSystemClassLoader addAllDirectLoadClasses(Collection<T> clss) {
        for (T obj : clss) {
            addDirectLoadClass(obj.getClass());
        }
        return this;
    }

    @Override
    protected Class<?> findClass(String name) throws ClassNotFoundException {
        byte[] classData = getClassData(name);
        if (classData == null) {
            throw new ClassNotFoundException();
        } else {
            return defineClass(name, classData, 0, classData.length);
        }
    }

    @Override
    protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
        Class<?> cls = findLoadedClass(name);
        if (cls == null) {// 该类没有加载
            // 如果该类需要由该加载器加载 ，则通过该加载器加载
            if (directLoadClassSet.contains(name)) {
                cls = findClass(name);
            } else {// 否则交给系统加载器加载
                cls = findSystemClass(name);
            }
        }
        if (cls == null) {
            throw new ClassNotFoundException(name);
        }
        if (resolve) {
            resolveClass(cls);
        }
        return cls;
    }

    /**
     * 获取指定的class文件数据
     * 
     * @param className
     * @return
     */
    private byte[] getClassData(String className) {
        try (InputStream ins = getClassInputStream(className); ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
            int bufferSize = 4096;
            byte[] buffer = new byte[bufferSize];
            int bytesNumRead = 0;
            while ((bytesNumRead = ins.read(buffer)) != -1) {
                baos.write(buffer, 0, bytesNumRead);
            }
            return baos.toByteArray();
        } catch (FileNotFoundException e) {
            LOGGER.error(className, e);
        } catch (IOException e) {
            LOGGER.error(className, e);
        }
        return null;
    }

    /**
     * 获取类文件输入流
     * 
     * @param className
     * @return
     * @throws IOException
     */
    private InputStream getClassInputStream(String className) throws IOException {
        String classFilePath = className.replace('.', File.separatorChar) + ".class";
        File file = new File(rootDir + File.separatorChar + classFilePath);
        if (file != null && file.exists()) {// 目录下存在类文件，优先读取类文件
            InputStream ins = new FileInputStream(file);
            return ins;
        } else {// 否则查看打包的脚本jar文件中是否存在类文件
            @SuppressWarnings("resource")
            JarFile jarFile = new JarFile(rootDir + File.separatorChar + scriptJarFileName);
            if (jarFile != null) {
                JarEntry jarEntry = jarFile.getJarEntry(classFilePath.replace(File.separatorChar, '/'));
                if (jarEntry != null) {
                    return jarFile.getInputStream(jarEntry);
                }
            }
        }
        throw new FileNotFoundException(className);
    }
}
